import matplotlib.pyplot as plt
import numpy as np

# 设置随机数种子
np.random.seed(5)

# ========= 准备数据 ============
# 直接采用np生成等差数列的方法，生成100个点，每个点的取值在-1~1之间
x_data = np.linspace(-1, 1, 100)
# y = 2x +1 + 噪声， 其中，噪声的维度与x_data一致
y_data = 2 * x_data + 1.0 + np.random.randn(*x_data.shape) * 0.4
# 画出随机生成数据的散点图
plt.scatter(x_data, y_data)
# 画出我们想要学习到的线性函数 y = 2x +1
plt.plot(x_data, 2 * x_data + 1.0, color='red', linewidth=3)
plt.show()

# ========= 构建模型 ============
import tensorflow as tf

# 定义模型函数
def model(x, w, b):
    return tf.multiply(x, w) + b


# 构建线性函数的斜率，变量w
w = tf.Variable(np.random.rand(), tf.float32, name="w0")
# 构建线性函数的截距，变量b
b = tf.Variable(0.0, name="b0")

# 定义损失函数
def loss(x, y, w, b):
    err = model(x, w, b) - y
    return tf.reduce_mean(tf.square(err))

# 迭代次数（训练轮数）
train_epochs = 10
# 学习率
learning_rate = 0.01
# 计算梯度
# 在 TensorFlow 2 中，使用 tf.GradientTape() 这一上下文管理器封装需要求导的计算步骤，并使用其 gradient() 方法求导
def grad(x, y, w, b):
    with tf.GradientTape() as tape:
        loss_ = loss(x, y, w, b)
    return tape.gradient(loss_, [w, b])


# ========== 训练模型（SGD，随机梯度下降） ===========
# 开始训练，轮数为 epoch，采用SGD随机梯度下降优化方法
step = 0   # 记录训练步数
display_step = 10  # 控制训练过程中，数据显示的频率
# 开始训练，轮数为 epoch，采用SGD随机梯度下降优化方法
for epoch in range(train_epochs):
    for xs, ys in zip(x_data, y_data):
        # 计算损失
        loss_ = loss(xs, ys, w, b)
        # 计算当前[w,b]点的梯度
        delta_w, delta_b = grad(xs, ys, w, b)
        # 更新[w,b]
        w.assign_sub(delta_w * learning_rate)
        b.assign_sub(delta_b * learning_rate)

        step += 1
        if step % display_step == 0:
            print('Training Epoch:', '%02d' % (epoch+1), "Step: %03d" % step, "loss=%.6f" % loss_)

    plt.plot(x_data, w.numpy() * x_data + b.numpy())  # 画图
plt.show()

# 打印结果
w_result = w.numpy()
b_result = b.numpy()
print("w: ", w_result)
print("b: ", b_result)

# 结果可视化
plt.scatter(x_data, y_data, label='Original Data')
plt.plot(x_data, x_data * w_result + b_result, label='Fitted Line')
plt.show()

# ======= 进行预测 ========
x_test = 3.21

predict = model(x_test, w_result, b_result)
print("预测值：%f" % predict)
